package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"server/config"

	"github.com/gin-gonic/gin"
)

// GenerateAESKey 生成指定长度的AES密钥（16/24/32字节对应AES-128/192/256）
func GenerateAESKey(keyLen int) (string, error) {
	if keyLen != 16 && keyLen != 24 && keyLen != 32 {
		return "", fmt.Errorf("AES密钥长度必须是16、24或32字节")
	}
	key := make([]byte, keyLen)
	_, err := io.ReadFull(rand.Reader, key)
	if err != nil {
		return "", fmt.Errorf("生成随机密钥失败: %v", err)
	}
	return base64.StdEncoding.EncodeToString(key), nil
}

// AESEncrypt 使用AES-CBC模式加密数据（接受Base64编码的密钥）
func AESEncrypt(plaintext []byte, keyBase64 string) (string, error) {
	// 解码Base64密钥
	key, err := base64.StdEncoding.DecodeString(keyBase64)
	if err != nil {
		return "", fmt.Errorf("Base64解码密钥失败: %v", err)
	}

	// 创建AES块
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建AES cipher失败: %v", err)
	}

	// 生成随机IV
	iv := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", fmt.Errorf("生成随机IV失败: %v", err)
	}

	// 填充明文到块大小的整数倍（PKCS#7填充）
	paddedPlaintext := pkcs7Padding(plaintext, block.BlockSize())

	// 创建CBC加密器
	mode := cipher.NewCBCEncrypter(block, iv)

	// 加密数据
	ciphertext := make([]byte, len(paddedPlaintext))
	mode.CryptBlocks(ciphertext, paddedPlaintext)

	// 拼接IV和密文
	combined := append(iv, ciphertext...)

	// 返回Base64编码的结果
	return base64.StdEncoding.EncodeToString(combined), nil
}

// AESDecrypt 使用AES-CBC模式解密数据（接受Base64编码的密钥）
func AESDecrypt(ciphertextBase64 string, keyBase64 string) ([]byte, error) {
	// 解码Base64密文
	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
	if err != nil {
		return nil, fmt.Errorf("Base64解码密文失败: %v", err)
	}

	// 解码Base64密钥
	key, err := base64.StdEncoding.DecodeString(keyBase64)
	if err != nil {
		return nil, fmt.Errorf("Base64解码密钥失败: %v", err)
	}

	// 检查密文长度是否合法
	if len(ciphertext) < aes.BlockSize {
		return nil, fmt.Errorf("密文长度过短")
	}

	// 提取IV
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	// 检查密文长度是否为块大小的整数倍
	if len(ciphertext)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("密文长度不是块大小的整数倍")
	}

	// 创建AES块
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("创建AES cipher失败: %v", err)
	}

	// 创建CBC解密器
	mode := cipher.NewCBCDecrypter(block, iv)

	// 解密数据
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	// 去除填充
	plaintext, err = pkcs7Unpadding(plaintext)
	if err != nil {
		return nil, fmt.Errorf("去除填充失败: %v", err)
	}

	return plaintext, nil
}

// pkcs7Padding 实现PKCS#7填充
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padtext...)
}

// pkcs7Unpadding 移除PKCS#7填充
func pkcs7Unpadding(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, fmt.Errorf("数据长度为0")
	}
	padding := int(data[length-1])
	if padding > length {
		return nil, fmt.Errorf("填充长度超过数据长度")
	}
	for i := length - padding; i < length; i++ {
		if data[i] != byte(padding) {
			return nil, fmt.Errorf("填充字节不一致")
		}
	}
	return data[:length-padding], nil
}

// DecryptMiddleware 解密请求体中间件
func DecryptMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 判断请求内容类型
		contentType := c.ContentType()
		if contentType != "application/x-www-form-urlencoded" {
			c.Next()
			return
		}
		// 读取原始请求体数据
		var buf bytes.Buffer
		_, err := buf.ReadFrom(c.Request.Body)
		if err != nil {
			// 处理读取错误，比如返回错误响应
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}
		// 将读取的请求体数据重新赋值给 c.Request.Body，以便后续处理能再次读取
		c.Request.Body = io.NopCloser(&buf)
		encryptedData := buf.String()

		// 获取 Base64 编码的 AES 密钥
		keyBase64 := config.Aes.Key

		// 调用解密函数解密
		decryptedData, err := AESDecrypt(encryptedData, keyBase64)
		if err != nil {
			// 解密失败处理，比如返回错误响应
			c.AbortWithError(http.StatusBadRequest, fmt.Errorf("解密请求数据失败: %v", err))
			return
		}

		// 将解密后的数据重新设置为请求体
		c.Request.Body = io.NopCloser(bytes.NewBuffer(decryptedData))

		// 继续执行后续的中间件或路由处理函数
		c.Next()
	}
}
